import pandas as pd
import numpy as np
import joblib
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, auc

# 假设你有 test_do() 函数返回 x_test, y_test
from transform_test import tes1_do  # ← 替换为你的实际模块名，比如 'transform_test'


def plot_roc_comparison():
    # 加载测试数据
    x_test, y_test = tes1_do()
    x_test = x_test.iloc[:, 0:20]

    # 加载两个模型
    rf_model = joblib.load('../model/tree.pkl')  # 随机森林
    xgb_model = joblib.load('../model/xgb.pkl')  # XGBoost

    # 获取预测概率（正类的概率）
    y_pred_rf = rf_model.predict_proba(x_test)[:, 1]
    y_pred_xgb = xgb_model.predict_proba(x_test)[:, 1]

    # 计算 ROC 曲线和 AUC
    fpr_rf, tpr_rf, _ = roc_curve(y_test, y_pred_rf)
    auc_rf = auc(fpr_rf, tpr_rf)

    fpr_xgb, tpr_xgb, _ = roc_curve(y_test, y_pred_xgb)
    auc_xgb = auc(fpr_xgb, tpr_xgb)

    # 绘图
    plt.figure(figsize=(8, 6))
    plt.plot(fpr_rf, tpr_rf, label=f'Random Forest (AUC = {auc_rf:.4f})', color='blue', linewidth=2)
    plt.plot(fpr_xgb, tpr_xgb, label=f'XGBoost (AUC = {auc_xgb:.4f})', color='red', linewidth=2)

    # 参考线（随机分类器）
    plt.plot([0, 1], [0, 1], 'k--', label='Random Classifier (AUC = 0.5)', linewidth=1)

    # 图形设置
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate (FPR)')
    plt.ylabel('True Positive Rate (TPR)')
    plt.title('ROC Curves: Random Forest vs XGBoost')
    plt.legend(loc='lower right')
    plt.grid(alpha=0.3)
    plt.tight_layout()

    # 保存或显示
    plt.savefig('../data/fig/roc_comparison.png', dpi=300)
    plt.show()


if __name__ == '__main__':
    plot_roc_comparison()
